import numpy as np
from PIL import Image
import os
import random
import shutil
import base64
from io import BytesIO


def remove_noise(img_np, threshold):
    """
    去噪音（达不到阈值的变为白色）
    :param img_np:
    :param threshold:
    :return:
    """
    is_sum = len(img_np.shape) >= 3
    height, width = img_np.shape[0], img_np.shape[1]
    for x in range(0, height):
        for y in range(0, width):
            if is_sum:
                if sum(img_np[x, y]) > threshold:
                    img_np[x, y] = 255
            else:
                if img_np[x, y] > threshold:
                    img_np[x, y] = 255
    return img_np


def cut_white_area(img_np, threshold):
    """
    切除白边
    :param img_np:
    :param threshold: 阈值
    :return:
    """
    is_sum = len(img_np.shape) >= 3
    height, width = img_np.shape[0], img_np.shape[1]
    minx = height
    miny = width
    maxx = 0
    maxy = 0
    for x in range(0, height):
        for y in range(0, width):

            is_not_white = False
            if is_sum and sum(img_np[x, y]) < threshold:
                is_not_white = True
            elif not is_sum and img_np[x, y] < threshold:
                is_not_white = True

            if is_not_white:
                if minx > x:
                    minx = x
                if miny > y:
                    miny = y
                if maxx < x:
                    maxx = x
                if maxy < y:
                    maxy = y
    # print("h:{} ,w:{}".format(maxx - minx, maxy - miny))
    new_img = img_np[minx:maxx + 1, miny:maxy + 1]
    return new_img



def resize(image_pil, width, height):
    """
    Resize PIL image keeping ratio and using white background.
    最后返回彩色
    """
    ratio_w = width / image_pil.width
    ratio_h = height / image_pil.height
    if ratio_w < ratio_h:
        # It must be fixed by width
        resize_width = width
        resize_height = round(ratio_w * image_pil.height)
    else:
        # Fixed by height
        resize_width = round(ratio_h * image_pil.width)
        resize_height = height
    image_resize = image_pil.resize((resize_width, resize_height), Image.ANTIALIAS)
    background = Image.new('RGBA', (width, height), (255, 255, 255, 255))
    offset = (round((width - resize_width) / 2), round((height - resize_height) / 2))
    background.paste(image_resize, offset)
    return background.convert('RGB')


def format_img_to_model(im, img_size):
    """
    为模型格式化参数
    :param img_size:
    :param im:
    :return:
    """
    captcha_array = np.array(im).astype('float32') / 255  # 向量化
    width, height = img_size
    captcha_array = captcha_array.reshape((width, height, -1))
    return captcha_array


def create_path(train_path, test_path, char_set):
    """
    创建文件夹
    :param char_set:
    :param train_path:
    :param test_path:
    :return:
    """
    for c_path in char_set:
        tra_path = train_path + '/' + c_path
        tes_path = test_path + '/' + c_path
        if not os.path.exists(tra_path):
            os.makedirs(tra_path)
        if not os.path.exists(tes_path):
            os.makedirs(tes_path)


def split_train_and_test(from_path, train_path, test_path, char_set):
    """
    切分测试集和验证集
    :param from_path:
    :param train_path:
    :param test_path:
    :param char_set:
    :return:
    """
    if not os.path.exists(from_path):
        os.makedirs(from_path)
    if not os.path.exists(train_path):
        os.makedirs(train_path)
    if not os.path.exists(test_path):
        os.makedirs(test_path)
    rate = 0.125
    # 创建分类文件夹
    create_path(train_path, test_path, char_set)

    fileList = os.listdir(from_path)
    testIndex = random.sample(range(len(fileList)), int(len(fileList) * rate))
    for i in range(len(fileList)):
        nextPath = train_path + '/' + fileList[i].split('-')[0]
        if i in testIndex:
            nextPath = test_path + '/' + fileList[i].split('-')[0]
        shutil.move(from_path + '/' + fileList[i], nextPath + '/' + fileList[i])
    print('切分训练与测试集完成')


def base64_to_image(base64_str, image_path=None, image_name=None):
    """
    base64转图片
    :param image_name:
    :param base64_str: base64编码
    :param image_path: 图片路径
    :return:
    """
    #     base64_data = re.sub('^data:image/.+;base64,', '', base64_str)
    byte_data = base64.b64decode(base64_str)
    image_data = BytesIO(byte_data)
    img = Image.open(image_data)
    if image_path and image_name:
        if not os.path.exists(image_path):
            os.makedirs(image_path)
        path = os.path.join(image_path, image_name)
        img.save(path)
    return img
